package cn.jbolt.ai.core.chain;

import cn.hutool.core.util.StrUtil;
import cn.jbolt.ai.app.config.AppConfigConstant;
import cn.jbolt.ai.app.entity.AiApp;
import cn.jbolt.ai.app.entity.AiAppMessage;
import cn.jbolt.ai.app.entity.AiAppSession;
import cn.jbolt.ai.app.vo.AiChatData;
import cn.jbolt.ai.app.vo.SocketData;
import cn.jbolt.ai.klb.entity.AiDoc;
import cn.jbolt.util.DateUtil;
import cn.jbolt.util.JSONUtil;
import com.fasterxml.jackson.databind.node.ArrayNode;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;

public class ChainContext  {


    /**
     * 请求消息
     */
    private AiAppMessage requestMsg;
    /**
     * 响应消息
     */
    private AiAppMessage responseMsg;

    private AiApp aiApp;

    private AiAppSession aiSession;

    private SocketData socketData;

    private WebSocketSession socketSession;

    private HashMap<String, Object> args = new HashMap<>();

    /**
     * 发送结束消息
     */
    public void sendOverMsg() {
        SocketData msg = SocketData.copy(socketData);
        msg.setAction(SocketData.ACTION_OVER);
        if (socketSession != null) {
            try {
                socketSession.sendMessage(new TextMessage(msg.toString()));
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }
    /**
     * 发送同步数据消息
     */
    public void sendSyncDataMsg() {
        SocketData msg = SocketData.copy(socketData);
        msg.setAction(SocketData.ACTION_SYNCDATA);
        msg.setData(JSONUtil.createObjectNode(
                "request", JSONUtil.createObjectNode("id",requestMsg.getId(), "createTime", DateUtil.format(requestMsg.getCreateTime())),
                "response", JSONUtil.createObjectNode("id",responseMsg.getId(), "createTime", DateUtil.format(responseMsg.getCreateTime()))));
        if (socketSession != null) {
            try {
                socketSession.sendMessage(new TextMessage(msg.toString()));
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }

    /**
     * 发送响应消息
     */
    public void sendResponseMsg(String msg,String format) {
        SocketData data = SocketData.copy(socketData);
        data.setAction(SocketData.ACTION_RESPONSE);
        AiChatData assistantMsg = AiChatData.createAssistantMsg(msg, format);
        data.setData(assistantMsg);
        if (socketSession != null) {
            try {
                socketSession.sendMessage(new TextMessage(data.toString()));
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }

    /**
     * 发送异常消息
     */
    public void sendErrorMsg(String errorMsg) {
        SocketData msg = SocketData.copy(socketData);
        msg.setAction(SocketData.ACTION_ERROR);
        msg.setData(new AiChatData().setMsg(errorMsg));
        if (socketSession != null) {
            try {
                socketSession.sendMessage(new TextMessage(msg.toString()));
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }

    /**
     * 发送参考消息
     * @param list
     */
    public void sendReference(List<AiDoc> list) {
        if (list == null || list.isEmpty()) return;
        ArrayNode array = JSONUtil.createArrayNode();
        for (AiDoc aiDoc : list) {
            array.add(aiDoc.convertToReference());
        }
        SocketData data = SocketData.copy(socketData);
        data.setAction(SocketData.ACTION_REFERENCE);
        data.setData(array);
        if (socketSession != null) {
            try {
                socketSession.sendMessage(new TextMessage(data.toString()));
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }


    public void sendSuccessMsgTip(String content) {
        sendResponseMsg( StrUtil.format("<msg-tip type='success'>{}</msg-tip>\n\n", content),"markdown");
    }
    public void sendInfoMsgTip(String content) {
        sendResponseMsg(StrUtil.format("<msg-tip type='info'>{}</msg-tip>\n\n", content),"markdown");
    }
    public void sendErrorMsgTip(String content) {
        sendResponseMsg( StrUtil.format("<msg-tip type='error'>{}</msg-tip>\n\n", content),"markdown");
    }



    public WebSocketSession getSocketSession() {
        return socketSession;
    }

    public void setSocketSession(WebSocketSession socketSession) {
        this.socketSession = socketSession;
    }

    public SocketData getSocketData() {
        return socketData;
    }

    public void setSocketData(SocketData socketData) {
        this.socketData = socketData;
    }

    public AiApp getAiApp() {
        return aiApp;
    }

    public void setAiApp(AiApp aiApp) {
        this.aiApp = aiApp;
    }

    public AiAppSession getAiSession() {
        return aiSession;
    }

    public void setAiSession(AiAppSession aiSession) {
        this.aiSession = aiSession;
    }

    /**
     * 获取原始提示词
     * @return
     */
    public String getOriPrompt(){
        return this.getAttr(AppConfigConstant.ATTR_KEY_PROMPT);
    }

    /**
     * 设置原始提示词
     * @param oriPrompt
     */
    public void setOriPrompt(String oriPrompt){
        this.args.put(AppConfigConstant.ATTR_KEY_PROMPT,oriPrompt);
    }


    public <T> T getAttr(String key) {
        if (StrUtil.isBlank(key)) {
            return null;
        }
        return (T) this.args.get(key);
    }
    public ChainContext setAttr(String key, Object value) {
        this.args.put(key, value);
        return this;
    }

    public HashMap<String, Object> getArgs() {
        return args;
    }

    public AiAppMessage getRequestMsg() {
        return requestMsg;
    }

    public void setRequestMsg(AiAppMessage requestMsg) {
        this.requestMsg = requestMsg;
    }

    public AiAppMessage getResponseMsg() {
        return responseMsg;
    }

    public void setResponseMsg(AiAppMessage responseMsg) {
        this.responseMsg = responseMsg;
    }
}
